

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader

import numpy as np
import scipy
import random
import threading
from typing import List, Dict, Tuple
from collections import Counter


# 实现简易的分词器
class WordTokenizer(object):
    def __init__(self):
        super(WordTokenizer, self).__init__()
    
    def split_words(self, content_str: str) -> List[str]:
        sentences_str = content_str.strip()
        sentences_str = sentences_str.lower()
        words_list = sentences_str.split(' ')
        
        return words_list  


# 封装语料预处理的逻辑
class CorpusData(object):
    def __init__(self, word_tokenizer: WordTokenizer, max_vocab_size: int):
        super(CorpusData, self).__init__()
        self.word_tokenizer = word_tokenizer
        self.max_vocab_size = max_vocab_size
        # 编码前的词汇序列
        self.original_words_list: List[str] = list()
        # 编码后的词汇序列
        self.encoded_words_list: List[str] = list()
        # 词汇与 id 的映射关系
        self.word2id_dict: Dict[str, int] = dict()
        self.id2word_dict: Dict[int, str] = dict()
        # 记录词汇表的列表
        self.vocab_words_list: List[str] = list()
        # 统计词汇表的频数
        self.vocab_counts_list: List[int] = list()
        
    def load_data(self, data_str: str) -> None:
        original_words_list = self.word_tokenizer.split_words(data_str)
        # 统计词汇出现的次数
        word2count_dict = Counter(original_words_list)
        # 只保留出现次数最多的前 N 个词汇，剩下的其它词汇用 <UNK> 代替
        word2count_dict = dict(word2count_dict.most_common(self.max_vocab_size - 1))
        word2count_dict['<UNK>'] = len(original_words_list) - sum(word2count_dict.values())
        
        vocab_words_list = sorted(word2count_dict, key=word2count_dict.get, reverse=True)
        vocab_counts_list = list()
        word2id_dict = dict()
        id2word_dict = dict()
        
        for index, word in enumerate(vocab_words_list):
            vocab_counts_list.append(word2count_dict[word])
            word2id_dict[word] = index
            id2word_dict[index] = word
            
        unknown_word_id = word2id_dict['<UNK>']
        # 编码词汇序列成对应的 id 序列
        encoded_words_list = [word2id_dict.get(word, unknown_word_id) for word in original_words_list]
        
        self.original_words_list = original_words_list
        self.encoded_words_list = encoded_words_list
        self.word2id_dict = word2id_dict
        self.id2word_dict = id2word_dict
        self.vocab_words_list = vocab_words_list
        self.vocab_counts_list = vocab_counts_list
        

# 自定义数据集
class SkipGramDataset(Dataset):
    def __init__(self, corpus_data: CorpusData, max_window_size: int, negative_sample_num: int):
        super(SkipGramDataset, self).__init__()
        corpus_words_list = corpus_data.encoded_words_list
        corpus_word_num = len(corpus_words_list)
        # 根据窗口大小忽略开头和结尾的 N 个词汇
        dataset_item_num = corpus_word_num - (2 * max_window_size)
        center_index_offset = max_window_size
        # 计算负样本的采样概率
        vocab_counts_list = corpus_data.vocab_counts_list
        vocab_counts_ndarray = np.array(vocab_counts_list, dtype=np.float32)
        negative_probs_ndarray = vocab_counts_ndarray ** 0.75
        negative_probs_ndarray = negative_probs_ndarray / np.sum(negative_probs_ndarray)
        
        self.dataset_item_num = dataset_item_num
        self.center_index_offset = center_index_offset
        self.max_window_size = max_window_size
        self.corpus_words_list = corpus_words_list
        self.corpus_words_tensor = torch.IntTensor(corpus_words_list)
        self.corpus_word_num = corpus_word_num
        # 每个正样本对应的负样本数目
        self.negative_sample_num = negative_sample_num
        self.negative_probs_tensor = torch.FloatTensor(negative_probs_ndarray)
        self.current_thread_data = threading.local()
        
    def __len__(self):
        return self.dataset_item_num
    
    def __getitem__(self, index) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        center_word_index = index + self.center_index_offset
        # 获取中心词
        center_word = self.corpus_words_list[center_word_index]
        center_words_tensor = torch.IntTensor([center_word])
        
        window_size = self.max_window_size
        start_index = center_word_index - window_size
        end_index = center_word_index + window_size + 1
        
        # 获取中心词附近的上下文词
        positive_words_list = [self.corpus_words_list[curr_index]  for curr_index 
                               in range(start_index, end_index) if curr_index != center_word_index]
        positive_words_tensor = torch.IntTensor(positive_words_list)
        
        # 使用线程本地变量去避免多线程情况下可能出现的逻辑错误
        try:
            negative_candidates_tensor = self.current_thread_data.negative_candidates_tensor
        except AttributeError:
            negative_candidates_tensor = self.negative_probs_tensor.clone()
            self.current_thread_data.negative_candidates_tensor = negative_candidates_tensor
            
        # 根据正样本计算要采样的负样本数目
        negative_word_num = self.negative_sample_num * len(positive_words_tensor)
        # 通过把概率置为 0 实现负样本采样时跳过正样本中的词汇
        negative_exclude_list = [self.corpus_words_list[curr_index]  for curr_index 
                                    in range(start_index, end_index)]
        negative_candidates_tensor[negative_exclude_list] = 0
        # 随机获取 N 个负样本
        negative_words_tensor = torch.multinomial(negative_candidates_tensor, negative_word_num, replacement=True)
        negative_candidates_tensor[negative_exclude_list] = self.negative_probs_tensor[negative_exclude_list]
        
        return center_words_tensor, positive_words_tensor, negative_words_tensor
            

# Skip-Gram模型
class SkipGramModel(nn.Module):
    def __init__(self, vocab_size: int, embedding_size: int):
        super(SkipGramModel, self).__init__()
        # 词汇表大小
        self.vocal_size = vocab_size
        # 词向量的维度大小
        self.embedding_size = embedding_size
        self.embedding = nn.Embedding(vocab_size, embedding_size)
        self.embedding_out = nn.Embedding(vocab_size, embedding_size)
        
    def forward(self, center_in: torch.Tensor, positive_out: torch.Tensor, negative_out: torch.Tensor) -> torch.Tensor:
        center_emb_in = self.embedding(center_in)
        # 转置该矩阵以满足矩阵相乘的前提条件，第二个矩阵的行数等于第一个矩阵的列数
        center_emb_in = torch.transpose(center_emb_in, dim0=2, dim1=1)
        positive_emb_out = self.embedding_out(positive_out)
        negative_emb_out = self.embedding_out(negative_out)
        
        # 正样本与中心词做矩阵乘法运算
        positive_prob = torch.bmm(positive_emb_out, center_emb_in)  
        positive_prob = torch.squeeze(positive_prob, dim=2)
        positive_prob = nn.functional.logsigmoid(positive_prob)
        positive_prob = torch.sum(positive_prob, dim=1)

        # 负样本与中心词做矩阵乘法运算
        negative_emb_out = torch.neg(negative_emb_out)
        negative_prob = torch.bmm(negative_emb_out, center_emb_in)
        negative_prob = torch.squeeze(negative_prob, dim=2)
        negative_prob = nn.functional.logsigmoid(negative_prob)
        negative_prob = torch.sum(negative_prob, dim=1)
        
        loss = positive_prob + negative_prob
        loss = torch.neg(loss)
        loss = torch.mean(loss)
        
        return loss


# 封装训练模型和测试模型的逻辑
class Word2vec(object):
    def __init__(self, corpus_data: CorpusData, embedding_size: int):
        super(Word2vec, self).__init__()
        self.corpus_data = corpus_data
        self.embedding_size = embedding_size
    
    def train_model(self, output_file_path: str, max_window_size: int, negative_sample_num: int,
                    epoch_num = 1, batch_size = 1, learning_rate = 0.1) -> None:
        """
        训练模型.
        
        Args:
          output_file_path: 训练结束后模型中的参数要保存的文件路径
          max_window_size: 最大窗口大小
          negative_sample_num: 每个正样本对应的负样本数目
          epoch_num: 迭代次数
          batch_size: 批次大小
          learning_rate: 学习率
        """
        
        # 此处通过设置随机方法的种子，可以复现训练的过程
        random.seed(1)
        np.random.seed(1)
        torch.manual_seed(1)
        
        if torch.cuda.is_available():
            device = 'cuda'
        else:
            device = 'cpu'
        
        dataset = SkipGramDataset(self.corpus_data, max_window_size=max_window_size, 
                                  negative_sample_num=negative_sample_num)
        dataloader = DataLoader(dataset, batch_size, shuffle=True)
        batch_num = len(dataloader)
        
        vocab_size = len(self.corpus_data.vocab_words_list)
        # 创建 Skip-Gram 模型
        model = SkipGramModel(vocab_size = vocab_size, embedding_size = self.embedding_size)
        model = model.to(device)
        optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
        model.train()
        
        for epoch_id in range(epoch_num):
            print(f'Epoch {epoch_id+1}/{epoch_num}')
            
            for batch_id, (center_in, positive_out, negative_out) in enumerate(dataloader):
                center_in = center_in.to(device)
                positive_out = positive_out.to(device)
                negative_out = negative_out.to(device)
                
                loss = model(center_in, positive_out, negative_out)
                # 根据误差反向计算梯度和更新权重
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
                
                if batch_id % 500 == 0:
                    print(f'loss: {loss.item():>10f}  [{batch_id+1:>5d}/{batch_num:>5d}]')
                    
        # 保存模型的所有参数
        state_dict = model.state_dict()        
        torch.save(state_dict, output_file_path)
    
    def test_model(self, model_state_path: str, test_words_list: List[str], nearest_word_num = 10) -> None:
        """
        测试模型.
        
        Args:
          model_state_path: 保存模型参数的文件路径
          test_words_list: 测试的词汇
          nearest_word_num: 每个测试词汇要获取最相似的词汇数
        """
        
        vocab_size = len(self.corpus_data.vocab_words_list)
        model = SkipGramModel(vocab_size = vocab_size, embedding_size = self.embedding_size)
        
        # 加载模型参数
        state_dict = torch.load(model_state_path)
        model.load_state_dict(state_dict)
        model.eval()
        
        word2id_dict = self.corpus_data.word2id_dict
        id2word_dict = self.corpus_data.id2word_dict
        # 模型中 embedding.weight 就是各个词汇对应的词向量
        weights_tensor = model.embedding.weight.detach()
        weights_ndarray = weights_tensor.numpy()
        
        for word in test_words_list:
            if (word not in word2id_dict):
                print(f'{word}: []')
                continue
            
            word_id = word2id_dict[word]
            word_weight = weights_ndarray[word_id]
            # 计算各个词汇与当前词汇的相似度
            similarity_list = [scipy.spatial.distance.cosine(curr_weight, word_weight) 
                              for curr_weight in weights_ndarray]
            similarity_ndarray = np.array(similarity_list)
            indexes_ndarray = similarity_ndarray.argsort()
            # 获取相似度最高的前 N 个词汇
            indexes_ndarray = indexes_ndarray[1:nearest_word_num+1]
            nearest_words_list =  [id2word_dict[curr_index] for curr_index in indexes_ndarray]

            print(f'{word}: {nearest_words_list}')
            

